import time

import numpy
from pymilvus import (
    connections,
    utility,
    FieldSchema,
    CollectionSchema,
    DataType,
    Collection,
)

fmt = "\n=== {:30} ===\n"
search_latency_fmt = "search latency = {:.4f}s"
num_entities, dim = 3000, 8

print(fmt.format("start connecting to Milvus"))
connections.connect("default", host="192.168.138.135", port="19530")

has = utility.has_collection("hello_milvus_test1")
print(f"Does collection hello_milvus exist in Milvus: {has}")

fields = [
    FieldSchema(name="pk", dtype=DataType.VARCHAR, is_primary=True, auto_id=False, max_length=100),
    FieldSchema(name="random", dtype=DataType.DOUBLE),
    FieldSchema(name="test", dtype=DataType.VARCHAR),
    FieldSchema(name="embeddings", dtype=DataType.FLOAT_VECTOR, dim=8)
]

schema = CollectionSchema(fields, "hello_milvus is the simplest demo to introduce the APIs")
print(fmt.format("Create collection `hello_milvus_test1_1`"))
# 操作哪张表；怎么查表
if has:
    hello_milvus = Collection("hello_milvus_test1_1")
else:
    hello_milvus = Collection("hello_milvus_test1_1", schema, consistency_level="Strong")
print(fmt.format("Start inserting entities"))


def insert(entities):
    insert_result = hello_milvus.insert(entities)

    hello_milvus.flush()
    print(f"Number of entities in Milvus: {hello_milvus.num_entities}")  # check the num_entites

    ################################################################################
    # 4. create index
    # We are going to create an IVF_FL  AT index for hello_milvus collection.
    # create_index() can only be applied to `FloatVector` and `BinaryVector` fields.
    print(fmt.format("Start Creating index IVF_FLAT"))
    index = {
        "index_type": "IVF_FLAT",
        "metric_type": "L2",
        "params": {"nlist": 128},
    }

    hello_milvus.create_index("embeddings", index)
    print(fmt.format("Start loading"))
    hello_milvus.load()